﻿
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#define __CUDACC__ 1
#include <mma.h>
using namespace nvcuda::wmma;

__global__ void computeWithoutPermute(float* A, float* B, float* C, unsigned int K) {
    fragment<matrix_a, 16, 16, 8, precision::tf32, row_major> frag_a;
    fragment<matrix_b, 16, 16, 8, precision::tf32, col_major> frag_b;
    fragment<accumulator, 16, 16, 8, float> frag_c = { };
    const unsigned int ldm = gridDim.x * 16;
    A += blockIdx.y * K * 16;
    B += blockIdx.x * K * 16;
    C += (blockIdx.y * ldm + blockIdx.x) * 16;
    for (unsigned int i = 0; i < K; i += 8) {
        load_matrix_sync(frag_a, A, K);
        load_matrix_sync(frag_b, B, K);
        mma_sync(frag_c, frag_a, frag_b, frag_c);
        A += 8;
        B += 8;
    }
    store_matrix_sync(C, frag_c, ldm, mem_row_major);
}

__global__ void computeAfterPermute(float* A, float* B, float* C, unsigned int K) {
    fragment<matrix_a, 16, 16, 8, precision::tf32, row_major> frag_a;
    fragment<matrix_b, 16, 16, 8, precision::tf32, col_major> frag_b;
    fragment<accumulator, 16, 16, 8, float> frag_c = { };
    const unsigned int ldm = gridDim.x * 16;
    A += blockIdx.y * K * 16;
    B += blockIdx.x * K * 16;
    C += (blockIdx.y * ldm + blockIdx.x) * 16;
    for (unsigned int i = 0; i < K; i += 8) {
#pragma unroll
        for (unsigned int j = 0; j < 4; ++j) {
            frag_a.x[j] = A[threadIdx.x];
            frag_b.x[j] = B[threadIdx.x];
            A += 32;
            B += 32;
        }
        mma_sync(frag_c, frag_a, frag_b, frag_c);
    }
    store_matrix_sync(C, frag_c, ldm, mem_row_major);
}

template<typename Use, typename Layout>
__global__ void permute(float* A, float* C, unsigned int K) {
    fragment<Use, 16, 16, 8, precision::tf32, Layout> frag;
    A += (blockIdx.y * K + blockIdx.x) * 16;
    C += (blockIdx.y * K + blockIdx.x * 16) * 16;
    load_matrix_sync(frag, A, K);
#pragma unroll
    for (unsigned int j = 0; j < 4; ++j) {
        C[threadIdx.x] = frag.x[j];
        C += 32;
    }
    A += 8;
    load_matrix_sync(frag, A, K);
#pragma unroll
    for (unsigned int j = 0; j < 4; ++j) {
        C[threadIdx.x] = frag.x[j];
        C += 32;
    }
}

#include <stdio.h>
#include <random>
#include <memory>

#define EPSILON 1e-1
template<typename T, unsigned int R, unsigned int C>
struct Matrix : public std::unique_ptr<T[]> {
    Matrix() : std::unique_ptr<T[]>(std::make_unique<T[]>(R* C)) {}

    T* operator[](size_t row) {
        return &(this->get())[row * C];
    }

    bool operator == (Matrix<T, R, C>& other) {
        bool res = true;
        for (int i = 0; i < R; ++i)
            for (int j = 0; j < C; ++j)
                if (abs((*this)[i][j] - other[i][j]) > EPSILON) {
                    printf("error at(%d, %d): %f %f\n", i, j, (*this)[i][j], other[i][j]);
                    res = false;
                }
        return res;
    }

    size_t memorySize() {
        return sizeof(T) * R * C;
    }
};

struct Event
{
    Event() {
        cudaEventCreate(&start);
        cudaEventRecord(start, 0);
    }

    ~Event() {
        cudaEventDestroy(start);
    }

    float elapsedTime() {
        cudaEvent_t stop;
        cudaEventCreate(&stop);
        cudaEventRecord(stop, 0);
        cudaEventSynchronize(stop);
        float time_ms;
        cudaEventElapsedTime(&time_ms, start, stop);
        cudaEventDestroy(stop);
        return time_ms * 1000;
    }
private:
    cudaEvent_t start;
};

int main(void)
{
#define M 2048
#define N 2048
#define K 4096
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dis(-1.0, 1.0);

    // k-major order matrix
    Matrix<float, M, K> a;
    Matrix<float, N, K> b;
    Matrix<float, M, N> c;
    Matrix<float, M, N> res;

    Matrix<float, M, K> a1;
    Matrix<float, N, K> b1;

    for (int i = 0; i < M; ++i)
        for (int j = 0; j < K; ++j)
            a[i][j] = dis(gen);

    for (int i = 0; i < N; ++i)
        for (int j = 0; j < K; ++j)
            b[i][j] = dis(gen);

    for (int i = 0; i < M; ++i)
        for (int j = 0; j < N; ++j)
        {
            float v = 0;
            for (int k = 0; k < K; ++k)
                v += a[i][k] * b[j][k];
            res[i][j] = v;
        }

    float* dev_a = nullptr;
    float* dev_b = nullptr;
    float* dev_c = nullptr;
    float* dev_a1 = nullptr;
    float* dev_b1 = nullptr;

    cudaMalloc((void**)&dev_a, a.memorySize());
    cudaMalloc((void**)&dev_b, b.memorySize());
    cudaMalloc((void**)&dev_c, c.memorySize());
    cudaMalloc((void**)&dev_a1, a.memorySize());
    cudaMalloc((void**)&dev_b1, b.memorySize());

    cudaMemcpy(dev_a, a.get(), a.memorySize(), cudaMemcpyHostToDevice);
    cudaMemcpy(dev_b, b.get(), b.memorySize(), cudaMemcpyHostToDevice);
    {
        Event event;
        computeWithoutPermute << <dim3(N / 16, M / 16), dim3(32) >> > (dev_a, dev_b, dev_c, K);
        printf("[computeWithoutPermute]\tcost time: %.1f us\n", event.elapsedTime());
    }
    cudaMemcpy(c.get(), dev_c, c.memorySize(), cudaMemcpyDeviceToHost);
    res == c;
    {
        Event event;
        permute<matrix_a, row_major> << <dim3(K / 16, M / 16), dim3(32) >> > (dev_a, dev_a1, K);
        permute<matrix_b, col_major> << <dim3(K / 16, N / 16), dim3(32) >> > (dev_b, dev_b1, K);
        computeAfterPermute << <dim3(N / 16, M / 16), dim3(32) >> > (dev_a1, dev_b1, dev_c, K);
        printf("[computeWithPermute]\tcost time: %.1f us\n", event.elapsedTime());
    }
    cudaMemcpy(c.get(), dev_c, c.memorySize(), cudaMemcpyDeviceToHost);
    res == c;
    {
        Event event;
        computeWithoutPermute << <dim3(N / 16, M / 16), dim3(32) >> > (dev_a, dev_b, dev_c, K);
        printf("[computeWithoutPermute]\tcost time: %.1f us\n", event.elapsedTime());
    }
    {
        Event event;
        permute<matrix_a, row_major> << <dim3(K / 16, M / 16), dim3(32) >> > (dev_a, dev_a1, K);
        permute<matrix_b, col_major> << <dim3(K / 16, N / 16), dim3(32) >> > (dev_b, dev_b1, K);
        computeAfterPermute << <dim3(N / 16, M / 16), dim3(32) >> > (dev_a1, dev_b1, dev_c, K);
        printf("[computeWithPermute]\tcost time: %.1f us\n", event.elapsedTime());
    }
    cudaFree(dev_a);
    cudaFree(dev_b);
    cudaFree(dev_c);
    cudaFree(dev_a1);
    cudaFree(dev_b1);
    cudaDeviceReset();
    return 0;
}